package wmi

import (
	"context"
	"log"
	"sync"

	"github.com/zzl/go-com/com"
	"github.com/zzl/go-wmi/wmi"
)

type ProcessMonitor struct {
	ctx     context.Context
	ctxFunc context.CancelFunc
	once    sync.Once

	onStart func(info *ProcessStartInfo)
	onStop  func(info *ProcessStopInfo)
}

// NewProcessMonitor 进程启动/退出监视
func NewProcessMonitor(onStart func(info *ProcessStartInfo), onStop func(info *ProcessStopInfo)) *ProcessMonitor {
	ctx, ctxFunc := context.WithCancel(context.TODO())
	return &ProcessMonitor{
		ctx:     ctx,
		ctxFunc: ctxFunc,
		onStart: onStart,
		onStop:  onStop,
	}
}

func (tis *ProcessMonitor) Open() error {
	tis.once.Do(func() {
		go tis.monitorStart()
		go tis.monitorStop()
	})
	return nil
}

func (tis *ProcessMonitor) Close() {
	tis.ctxFunc()
}

// 监视进程的启动 Win32_ProcessStartTrace
func (tis *ProcessMonitor) monitorStart() {
	const wql = "SELECT * FROM Win32_ProcessStartTrace"

	com.Initialize()
	defer com.Uninitialize()
	defer com.NewScope().Leave()

	service := wmi.ConnectDefault()

	eventSource := service.ExecNotificationQuery(wql)
	if eventSource == nil {
		// if the error is "access denied",
		// try running with administrative privileges.
		log.Printf("ExecNotificationQuery fail. %v", com.GetLastErrorInfo().String())
		return
	}

	for {
		select {
		case <-tis.ctx.Done():
			return
		default:
		}

		/*
			ParentProcessID = 8260;
			ProcessID = 4204;
			ProcessName = "conhost.exe";
			SessionID = 1;
			Sid = {1, 5, 0, 0, 0, 0, 0, 5, 21, 0, 0, 0, 210, 54, 214, 99, 215, 241, 231, 4, 164, 76, 153, 81, 233, 3, 0, 0};
			TIME_CREATED = "133398434657259102";
		*/

		com.WithScope(func() {
			event := eventSource.NextEvent()

			ppid := event.Get("ParentProcessID").Int()
			name := event.Get("ProcessName").String()
			pid := event.Get("ProcessID").Int()
			// fmt.Printf("%s(%d) ppid:%v Started.\n", name, pid, ppid)

			if tis.onStart != nil {
				tis.onStart(&ProcessStartInfo{
					ParentProcessID: ppid,
					ProcessName:     name,
					ProcessID:       pid,
				})
			}
		})
	}
}

// 监视进程的关闭 Win32_ProcessStopTrace
func (tis *ProcessMonitor) monitorStop() {
	const wql = "SELECT * FROM Win32_ProcessStopTrace"

	com.Initialize()
	defer com.Uninitialize()
	defer com.NewScope().Leave()

	service := wmi.ConnectDefault()

	eventSource := service.ExecNotificationQuery(wql)
	if eventSource == nil {
		log.Printf("ExecNotificationQuery fail. %v", com.GetLastErrorInfo().String())
		return
	}

	for {
		select {
		case <-tis.ctx.Done():
			return
		default:
		}

		/*
			ExitStatus = 1;
			ParentProcessID = 0;
			ProcessID = 20312;
			ProcessName = "CalculatorApp.";
			SessionID = 0;
			Sid = {1, 5, 0, 0, 0, 0, 0, 5, 21, 0, 0, 0, 210, 54, 214, 99, 215, 241, 231, 4, 164, 76, 153, 81, 233, 3, 0, 0};
			TIME_CREATED = "133398436292781311";
		*/

		com.WithScope(func() {
			event := eventSource.NextEvent()

			ppid := event.Get("ParentProcessID").Int()
			name := event.Get("ProcessName").String()
			pid := event.Get("ProcessID").Int()
			exitCode := event.Get("ExitStatus").Int()
			// fmt.Printf("%s(%d) Exited: %v.\n", name, pid, exitCode)

			if tis.onStop != nil {
				tis.onStop(&ProcessStopInfo{
					ParentProcessID: ppid,
					ProcessName:     name,
					ProcessID:       pid,
					ExitStatus:      exitCode,
				})
			}
		})
	}
}

type ProcessStartInfo struct {
	ParentProcessID int
	ProcessName     string
	ProcessID       int
}

type ProcessStopInfo struct {
	ParentProcessID int
	ProcessName     string
	ProcessID       int
	ExitStatus      int
}
